import json
from loguru import logger
from root_agent.llm import LLMInterface
from mem0 import MemoryClient

class ContextManager:
    def __init__(self):
        # Load MEM0_API_KEY from settings.txt
        mem0_api_key = None
        settings = self._load_settings()
        mem0_api_key = settings.get("MEM0_API_KEY")
        
        self.mem0 = MemoryClient(api_key=mem0_api_key)
        self.chat_history_ids = [] # To keep track of chat history entry IDs in mem0
        self.user_id = settings.get("userid", "test111")
        self.open_id = settings.get("openid", "zyinfoai")
        self.current_task_focus = {
            "work_dir": None,
            "task_id": None,
            "project_name": None
        }
        self.environment_snapshot = {}
        self.user_profile_summary = {}
        self.llm = LLMInterface()
        self.max_history_length = 10 # Summarize history if it exceeds 10 entries
        self.max_history_size_bytes = 20 * 1024 # 20KB

    def _load_settings(self):
        settings = {}
        try:
            with open("./settings.txt", 'r', encoding='utf-8') as f:
                settings = json.load(f)
        except FileNotFoundError:
            logger.error("Error: settings.txt not found. Cannot load MEM0_API_KEY.")
        except json.JSONDecodeError:
            logger.error("Error: Could not decode JSON from settings.txt.")
        return settings

    def update_history(self, role, message):
        # Store each message as a separate memory entry in mem0
        # We use a simple user_id for now, but this can be made dynamic
        user_id = self.user_id 
        memory_entry = [{"role": role, "content": message}]
        add_params = {"user_id": user_id, "metadata": {"type": "chat_history"}, "agent_id": self.open_id}
        logger.debug(f"Attempting mem0.add with entry: {memory_entry} and params: {add_params}")
        entry = self.mem0.add(memory_entry, **add_params)
        logger.debug(f"Mem0 add entry response: {entry}") # Debugging line
        if entry and isinstance(entry, dict) and 'results' in entry and \
           isinstance(entry['results'], list) and len(entry['results']) > 0 and \
           isinstance(entry['results'][0], dict) and 'id' in entry['results'][0]:
            self.chat_history_ids.append(entry['results'][0]['id']) # Store the ID of the new entry
        else:
            logger.warning(f"Mem0 add operation did not return a valid ID or expected structure: {entry}")

        # Check if history needs summarization
        self._summarize_history()

    def _summarize_history(self):
        user_id = self.user_id # Consistent user_id

        # Retrieve all chat history entries from mem0
        all_history_entries = self.mem0.get_all(user_id=user_id, metadata={"type": "chat_history"}, agent_id=self.open_id)
        if not isinstance(all_history_entries, list):
            all_history_entries = [] # Ensure it's a list even if mem0 returns something unexpected
        
        # Filter out entries that don't have 'created_at' or are not dictionaries
        all_history_entries = [entry for entry in all_history_entries if isinstance(entry, dict) and 'created_at' in entry]

        # Sort entries by creation time to ensure correct order for summarization
        all_history_entries.sort(key=lambda entry: entry['created_at'])

        # Convert mem0 entries to the format expected by the existing summarization logic
        current_history_for_summarization = []
        for entry in all_history_entries:
            if entry and isinstance(entry, dict) and entry.get('memory') and entry.get('id'):
                role = "unknown"
                if entry.get('metadata') and isinstance(entry.get('metadata'), dict):
                    role = entry['metadata'].get("role", "unknown")
                current_history_for_summarization.append({"role": role, "content": entry['memory'], "id": entry['id']})
            else:
                logger.warning(f"Skipping malformed history entry: {entry}")

        # Calculate current size of the history
        current_history_size = sum(len(json.dumps(entry)) for entry in current_history_for_summarization)

        # Only summarize if history exceeds the max size
        if current_history_size <= self.max_history_size_bytes:
            return

        # Determine how many entries to summarize (e.g., summarize the oldest half)
        num_entries_to_summarize = len(current_history_for_summarization) // 2
        if num_entries_to_summarize == 0: # Avoid infinite loop if history is too small to split
            return

        history_to_summarize = current_history_for_summarization[:num_entries_to_summarize]
        remaining_history = current_history_for_summarization[num_entries_to_summarize:]

        # Format history for LLM summarization
        formatted_history = "\n".join([f"{entry["role"]}: {entry["content"]}" for entry in history_to_summarize])

        prompt = f"""Summarize the following conversation history concisely. Focus on key information, decisions, and outcomes. Preserve important details that might be relevant for future context.

Conversation History:
{formatted_history}

Concise Summary:"""
        summary_text = self.llm.generate_content(prompt, task_type="summarization").strip()

        # Delete the summarized entries from mem0
        ids_to_delete = [entry['id'] for entry in history_to_summarize]
        for entry_id in ids_to_delete:
            self.mem0.delete(entry_id, user_id=user_id, agent_id=self.open_id)
        
        # Add the summary back to mem0 as a new entry
        summary_entry_response = self.mem0.add([{"role": "system", "content": f"SUMMARY OF PAST CONVERSATION: {summary_text}"}], user_id=user_id, metadata={"role": "system", "type": "chat_history_summary"})
        if summary_entry_response and isinstance(summary_entry_response, dict) and 'results' in summary_entry_response and \
           isinstance(summary_entry_response['results'], list) and len(summary_entry_response['results']) > 0 and \
           isinstance(summary_entry_response['results'][0], dict) and 'id' in summary_entry_response['results'][0]:
            summary_entry_id = summary_entry_response['results'][0]['id']
            # Update chat_history_ids to reflect the changes (remove deleted, add new summary ID)
            self.chat_history_ids = [entry_id for entry_id in self.chat_history_ids if entry_id not in ids_to_delete]
            self.chat_history_ids.append(summary_entry_id)
        else:
            logger.warning(f"Mem0 summary add operation did not return a valid ID or expected structure: {summary_entry_response}")

    def update_task_focus(self, work_dir=None, task_id=None, project_name=None):
        if work_dir:
            self.current_task_focus["work_dir"] = work_dir
        if task_id:
            self.current_task_focus["task_id"] = task_id
        if project_name:
            self.current_task_focus["project_name"] = project_name

    def update_environment_snapshot(self, snapshot_data):
        self.environment_snapshot.update(snapshot_data)

    def update_user_profile_summary(self, profile_data):
        self.user_profile_summary.update(profile_data)

    def get_context(self):
        user_id = self.user_id # Consistent user_id

        # Retrieve all chat history entries (including summaries) from mem0
        all_history_entries = self.mem0.get_all(user_id=user_id, metadata={"type": "chat_history"}, agent_id=self.open_id)
        if not isinstance(all_history_entries, list):
            all_history_entries = []
        summary_entries = self.mem0.get_all(user_id=user_id, metadata={"type": "chat_history_summary"}, agent_id=self.open_id)
        if not isinstance(summary_entries, list):
            summary_entries = []

        # Combine and filter out malformed entries, then sort by creation time
        combined_history = []
        for entry in all_history_entries + summary_entries:
            if isinstance(entry, dict) and 'created_at' in entry:
                combined_history.append(entry)
            else:
                logger.warning(f"Skipping malformed combined history entry: {entry}")
        combined_history.sort(key=lambda entry: entry['created_at'])

        # Format for the rest of the system
        formatted_chat_history = []
        for entry in combined_history:
            if entry and isinstance(entry, dict) and entry.get('memory'):
                role = "unknown"
                if entry.get('metadata') and isinstance(entry.get('metadata'), dict):
                    role = entry['metadata'].get("role", "unknown")
                formatted_chat_history.append({"role": role, "content": entry['memory']})
            else:
                logger.warning(f"Skipping malformed formatted chat history entry: {entry}")

        return {
            "history": formatted_chat_history,
            "task_focus": self.current_task_focus,
            "environment": self.environment_snapshot,
            "user_profile": self.user_profile_summary
        }
